import os.path
import json
import time
from utils import convert_dir_utf8, convert_txt_2_json, post_book, post_page
from meilisearch import Client
from qdrant_client import QdrantClient
from qdrant_client.http import models
from sentence_transformers import SentenceTransformer

MEILISEARCH_URI = "http://192.168.2.230:7700"
MEILISEARCH_MASTERKEY = "janl-key"

meili_client = Client(MEILISEARCH_URI, MEILISEARCH_MASTERKEY)
client = QdrantClient(host="192.168.2.230", port=6333)
model = SentenceTransformer("BAAI/bge-small-zh-v1.5")
def convert_dir_file_encode_2_utf8(sdir, ddir):
    """
    将目录中的文件编码格式，由非utf-8转为utf-8编码
    :param sdir: 源文件目录
    :param ddir: 目标文件目录
    :return:
    """
    print("-" * 72)
    print("start convent txt file from gb2312 to utf-8")

    convert_dir_utf8(sdir, ddir)
    print("-" * 72)

def convert_txt_dir_2_json(txt_dir, json_dir="json"):
    """
    将文本文件转换为json格式，用于构建meilisearch索引
    :param txt_dir:  文本文件目录
    :param json_dir:  json文件目录
    :return:
    """
    # curr_index = meili_client.index("books")
    # curr_index.delete_all_documents()
    books = []
    bookid = 1
    s_dir = txt_dir
    d_dir = json_dir
    if json_dir == "json":
        json_dir = os.path.join(txt_dir, json_dir)
    if not os.path.exists(json_dir):
        os.mkdir(json_dir)

    for file in os.listdir(s_dir):
        if os.path.isfile(full_name := os.path.join(s_dir, file)):
            try:
                book = convert_txt_2_json(bookid, full_name, json_dir)
                books.append(book)
                bookid = bookid + 1
            except Exception as e:
                print(f"convert json:{full_name}, encode error:{e}")
        # else:
        #     if file == "json":
        #         continue
        #     else:
        #         tmp_sdir = os.path.join(s_dir, file)
        #         tmp_ddir = os.path.join(json_dir, file)
        #         if not os.path.exists(tmp_ddir):
        #             os.mkdir(tmp_ddir)
        #         convert_txt_2_json(bookid, tmp_sdir, tmp_ddir)

        if bookid > 647:
            break

    return books



def update_meilisearch_index(start_bookid, index_name, json_path):
    curr_index = meili_client.index(index_name)
    print("create meilisearch meilisearch:", curr_index.uid)
    # documents = [
    #     {"id": 1, "title": "Inception", "genre": "Action"},
    #     {"id": 2, "title": "Interstellar", "genre": "Sci-Fi"}
    # ]
    books = []
    book = {}
    docs = []
    bookid = start_bookid
    flag = False
    bookname = ""
    while True:
        docs = []
        for file in os.listdir(json_path):
            if os.path.isfile(full_name := os.path.join(json_path, file)):
                try:
                    if file.startswith("book_"+str(bookid)+"_"):
                        with open(os.path.join(json_path, file), 'r', encoding='utf-8') as f:
                            data_dict = json.load(f)
                            if flag == False:
                                bookname = data_dict["book_name"]
                                book_sheets = data_dict["book_sheets"]
                                flag = True
                            docs.append(data_dict)
                            # print("docs len:", len(docs))
                except Exception as e:
                    print(f"convert json:{full_name}, encode error:{e}")

        book = {}
        book["id"] = bookid
        book["name"] = bookname
        book["docs"] = docs
        # r = input(f"确认，将《{bookname}》共{str(book_sheets)}篇内容构建成为索引文档吗？Yes/No/Cancal: ")
        r = "Y"
        if r.upper().strip() == "N":
            print("ask result is No.")
            book["update"] = False
            books.append(book)
            bookid = bookid + 1
            flag = False
            continue
        elif r.upper().strip() == "Y":
            print("ask result is Yes.")
            book["update"] = True
            books.append(book)
            documents = book["docs"]
            print("bookid:", bookid, "bookname:", bookname,  "documents len:", len(documents), " update:", book["update"])
            if len(documents) > 0 and book["update"]:
                # print("begin add meilisearch documents:", len(documents))
                # curr_index.add_documents(documents, "id")
                for doc in documents:
                    print(doc)
            bookid = bookid + 1
            flag = False
        else:
            print("ask result is Cancle.")
            break

        if bookid > 1:
            break


def splite_content_len512(books):
    sort = 1
    title_count = 1
    new_pages = []
    for book in books:
        # print(book)
        if book["sheets"] > 2:
            sort = 1
            new_pages = []
            # print(book["name"])
            # print(book["sort"], "book.id:", book["id"], "book.name:", book["name"], "book.sheets:", book["sheets"])
            # 拆分文章片段，使文章片段<512字。
            for page in book["pages"]:
                if len(page["content"]) < 512:
                    title_count = 1
                    np = {}
                    np["id"] = 0
                    np["book_id"] = 0
                    np["book_name"] = page["book_name"]
                    np["title"] = page["title"]
                    np["content"] = page["content"]
                    np["status"] = 1
                    np["sort"] = sort
                    new_pages.append(np)
                    sort = sort + 1
                else:
                    title_count = 1
                    lines = page["content"].split("\n")
                    sc = ""
                    for line in lines:
                        if len(sc) + len(line) < 512:
                            sc = sc + line + "\n"
                        else:
                            np = {}
                            np["id"] = 0
                            np["book_id"] = 0
                            np["book_name"] = page["book_name"]
                            np["title"] = page["title"] + "-" + str(title_count)
                            np["content"] = sc
                            np["status"] = 1
                            np["sort"] = sort
                            new_pages.append(np)
                            sort = sort + 1
                            title_count = title_count + 1
                            sc = ""
                    np = {}
                    np["id"] = 0
                    np["book_id"] = 0
                    np["book_name"] = page["book_name"]
                    np["title"] = page["title"] + "-" + str(title_count)
                    np["content"] = sc
                    np["status"] = 1
                    np["sort"] = sort
                    new_pages.append(np)
            book["pages"] = new_pages
            book["sheets"] = len(new_pages)
            # print(book["sort"], "book.id:", book["id"], "book.name:", book["name"], "book.sheets:", book["sheets"])
            # print(book["sort"], "----->" + str(page["sort"]) + " " + page["title"] + "---" + str(len(page["content"])))

    return books

def save_page_to_qdrant(v_points):
    if not client.collection_exists("chinese_medicine"):
        client.create_collection("chinese_medicine",
                vectors_config=models.VectorParams(size=512, distance=models.Distance.COSINE))

    # print("id.type:", type(page_id))
    # print("book.id:", book_id, "book.name:", page["book_name"], "page.id:", page_id, "page.title:", page["title"], "page.len:", len(content))
    # payload = {"pid": page_id, "bid": book_id, "name": book_name, "title": title, "content": content}
    client.upsert(
        collection_name="chinese_medicine",
        points=v_points
    )

def build_points(pages):
    v_points = []
    for page in pages:
        embeddings = model.encode(page["content"])
        book_id = int(page["book_id"])
        page_id = int(page["id"])
        content = page["content"]
        title = page["title"]
        sort = int(page["sort"])
        book_name = page["book_name"]
        payload = {"pid": page_id, "bid": book_id, "name": book_name, "title": title, "sort": sort, "content": content}
        point = models.PointStruct(id=page_id, payload=payload, vector=embeddings)
        v_points.append(point)
    return v_points

def save_books(books):
    book_id = 1
    page_id = 1
    for book in books:
        book["id"] = 0
        book_id = post_book(book_data=book)
        book["sort"] = book_id
        v_pages = []
        for page in book["pages"]:
            page["book_id"] = book_id
            page_id = post_page(page)
            # print("save page.id:", page_id)
            page["id"] = page_id
            v_pages.append(page)
            # save_page_to_qdrant(page)
            # page_id = page_id + 1
            # time.sleep(0.5)
        v_points = build_points(v_pages)
        print("all pages:", len(v_points))
        if len(v_points) > 1000:
            groups = [v_points[i:i+1000] for i in range(0, len(v_points), 1000)]
            for v in groups:
                save_page_to_qdrant(v)
                print("save pages:", len(v))
        else:
            save_page_to_qdrant(v_points)
            print("save pages:", len(v_points))

        # book_id = book_id + 1

def main():
    source_dir = r"E:\ali-disk\中医古籍700本"
    dest_dir = r"D:\BaiduNetdiskDownload\中医书\书"
    json_dir = r"D:\BaiduNetdiskDownload\中医书\书\json"
    print("-" * 36)
    print("|" + "中医书文本预处理程序".center(34) + "|")
    print("-" * 36)
    # if input_yes_no("需要做文件编码转换吗？ Yes/No?: "):
    #     print("step 1: 将文本文件编码转换为utf-8编码...")
    #     source_dir = input_path("请指定源文件路径或回车", source_dir)
    #     print("your source path:", source_dir)
    #     dest_dir = input_path("请指定目标文件路径或回车", dest_dir)
    #     print("your dest path:", dest_dir)
    #     convert_dir_file_encode_2_utf8(source_dir, dest_dir)
    # if input_yes_no("需要将文本文档转换为josn格式吗？ Yes/No?: "):
    #     print("step 2: 将文本文件转换为Json格式...")
    #     json_dir = input_path("请指定Json文件路径或回车", json_dir)
    books = convert_txt_dir_2_json(dest_dir, json_dir)
    # if input_yes_no("需要将json文件构建到索引吗？ Yes/No?: "):
    #     print("step 3: 使用Json文件构建Meilisearch索引...")
    #     json_path = os.path.join(dest_dir, json_dir)
    #     json_path = input_path("请指定Json文件路径或回车", json_path)
    #     update_meilisearch_index(1, "books", json_path)
    new_books = splite_content_len512(books)
    save_books(new_books)

if __name__ == "__main__":
    main()